# mongo_utils.py
import pymongo
import logging
from typing import List, Dict, Optional, Union
from contextlib import contextmanager

logger = logging.getLogger(__name__)


class MongoDBClient:
    """
    MongoDB工具类，支持连接、增删改查等操作
    Usage:
        client = MongoDBClient("user", "pwd", "127.0.0.1", 27017, "admin", "my_db")

        # 使用上下文管理
        with client.session() as db:
            db.insert("users", {"name": "Alice"})
            result = db.find("users", {"name": "Alice"}, one=True)
            print(result)
    """

    def __init__(
            self,
            username: str,
            password: str,
            host: str,
            port: Union[str, int],
            auth_db: str,
            db_name: str
    ) -> None:
        self.username = username
        self.password = password
        self.host = host
        self.port = port
        self.auth_db = auth_db
        self.db_name = db_name
        self._client: Optional[pymongo.MongoClient] = None
        self._db: Optional[pymongo.database.Database] = None

    def connect(self) -> None:
        """建立连接（延迟初始化）"""
        if not self._client:
            uri = (
                f"mongodb://{self.username}:{self.password}@{self.host}:{self.port}/"
                f"?authSource={self.auth_db}&readPreference=primary&appname=MongoDB%20Compass&ssl=false"
            )
            self._client = pymongo.MongoClient(uri)
            self._db = self._client[self.db_name]
            logger.info(f"Connected to MongoDB: {self.host}:{self.port}/{self.db_name}")

    def get_db(self):
        self.connect()
        return self._db

    def get_client(self):
        self.connect()
        return self._client

    def close(self):
        if self._client:
            self._client.close()
            logger.info("MongoDB connection closed")
            self._client = None
            self._db = None

    @contextmanager
    def session(self):
        """上下文管理连接"""
        self.connect()
        try:
            yield self
        finally:
            self.close()

    def find(self, table: str, query: Dict, projection: Optional[Dict] = None, one: bool = False):
        """执行 find 或 find_one 查询"""
        self.connect()
        collection = self._db[table]
        if one:
            return collection.find_one(query, projection)
        return collection.find(query, projection)

    def insert(self, table: str, data: Union[Dict, List[Dict]], many: bool = False):
        """插入文档"""
        self.connect()
        collection = self._db[table]
        if many:
            return collection.insert_many(data)
        return collection.insert_one(data)

    def update(self, table: str, filter_query: Dict, update_data: Dict, many: bool = False):
        """更新文档"""
        self.connect()
        collection = self._db[table]
        if many:
            return collection.update_many(filter_query, update_data)
        return collection.update_one(filter_query, update_data)

    def delete(self, table: str, query: Dict, many: bool = False):
        """删除文档"""
        self.connect()
        collection = self._db[table]
        if many:
            return collection.delete_many(query)
        return collection.delete_one(query)
